import numpy as np

from src.gait_scheduler.ramp_generator import RampGenerator
from ..robot_state import RobotState


class Joystick:
    def __init__(self, dt: float) -> None:
        super().__init__()
        self.dt = dt
        self.vx_L_gen = RampGenerator(dt)
        self.vy_L_gen = RampGenerator(dt)
        self.wz_L_gen = RampGenerator(dt)
        self.theta_z_gen = RampGenerator(dt)

        self.thetaZ = 0.0
        self.vx_W = 0.0
        self.vy_W = 0.0
        self.px_W = 0.0
        self.py_W = 0.0
        self.vx_L = 0.0
        self.vy_L = 0.0
        self.wz_L = 0.0

    def set_vx_des_L_para(self, vx_des_L: float, time: float):
        self.vx_L_gen.set_para(vx_des_L, time)

    def set_vy_des_L_para(self, vy_des_L: float, time: float):
        self.vy_L_gen.set_para(vy_des_L, time)

    def set_wz_des_L_para(self, wz_des_L: float, time: float):
        self.wz_L_gen.set_para(wz_des_L, time)

    def setIniPos(self, pos_x: float, pos_y: float, theta_z: float):
        self.px_W = pos_x
        self.py_W = pos_y
        self.thetaZ = theta_z

    def step(self):
        self.vx_L = self.vx_L_gen.step()
        self.vy_L = self.vy_L_gen.step()
        self.wz_L = self.wz_L_gen.step()
        self.thetaZ += self.wz_L * self.dt
        self.vx_W = np.cos(self.thetaZ) * self.vx_L - np.sin(self.thetaZ) * self.vy_L
        self.vy_W = np.sin(self.thetaZ) * self.vx_L + np.cos(self.thetaZ) * self.vy_L
        self.px_W += self.vx_W * self.dt
        self.py_W += self.vy_W * self.dt

    def write_state(self, robot_state: RobotState):
        robot_state.js_pos_des[0] = self.px_W
        robot_state.js_pos_des[1] = self.py_W
        robot_state.js_vel_des[0] = self.vx_W
        robot_state.js_vel_des[1] = self.vy_W
        robot_state.js_eul_des[2] = self.thetaZ
        robot_state.js_omega_des[2] = self.wz_L

    def reset(self):
        self.vx_L_gen.reset_out(0.0)
        self.vy_L_gen.reset_out(0.0)
        self.wz_L_gen.reset_out(0.0)
        self.vx_L = 0.0
        self.vy_L = 0.0
        self.wz_L = 0.0
        self.thetaZ = 0.0
